import os
import torch
import torch.nn as nn
import torch.optim as optim
from torchvision import datasets, transforms, utils
from torch.utils.data import DataLoader
import numpy as np
import matplotlib.pyplot as plt

# 1. 数据准备
def prepare_data():
    transform = transforms.Compose([
        transforms.ToTensor(),
        transforms.Normalize((0.5,), (0.5,))
    ])
    train_dataset = datasets.MNIST(root='./data', train=True, download=True, transform=transform)
    return DataLoader(train_dataset, batch_size=128, shuffle=True)

# 2. UNet模型
class UNet(nn.Module):
    def __init__(self):
        super().__init__()
        self.time_embed = nn.Sequential(
            nn.Embedding(1000, 32),
            nn.Linear(32, 784),
            nn.SiLU()
        )
        self.model = nn.Sequential(
            nn.Linear(784*2, 1024),
            nn.SiLU(),
            nn.Linear(1024, 1024),
            nn.SiLU(),
            nn.Linear(1024, 784)
        )
    
    def forward(self, x, t):
        B, C, H, W = x.shape
        x = x.view(B, -1)
        t_emb = self.time_embed(t)
        x = torch.cat([x, t_emb], dim=-1)
        return self.model(x).view(B, C, H, W)

# 3. 扩散模型
class Diffusion:
    def __init__(self, T=1000, device='cpu'):
        self.T = T
        self.device = device
        self.betas = torch.linspace(1e-4, 0.02, T).to(device)
        self.alphas = (1 - self.betas).to(device)
        self.alpha_bars = torch.cumprod(self.alphas, dim=0).to(device)
    
    def forward_diffusion(self, x0, t):
        # Ensure t is on the same device as alpha_bars
        t = t.to(self.alpha_bars.device)
        noise = torch.randn_like(x0)
        alpha_bar = self.alpha_bars[t].view(-1, 1, 1, 1)
        xt = torch.sqrt(alpha_bar) * x0 + torch.sqrt(1 - alpha_bar) * noise
        return xt, noise
    
    def train_step(self, model, x0, optimizer):
        t = torch.randint(0, self.T, (x0.shape[0],), device=x0.device)
        xt, noise = self.forward_diffusion(x0, t)
        pred_noise = model(xt, t)
        loss = nn.functional.mse_loss(pred_noise, noise)
        
        optimizer.zero_grad()
        loss.backward()
        optimizer.step()
        return loss.item()
    
    @torch.no_grad()
    def sample(self, model, n_samples=16, device='cpu'):
        model.eval()
        x = torch.randn((n_samples, 1, 28, 28), device=device)
        
        for t in reversed(range(self.T)):
            z = torch.randn_like(x) if t > 1 else 0
            beta_t = self.betas[t].to(device)
            alpha_t = self.alphas[t].to(device)
            alpha_bar_t = self.alpha_bars[t].to(device)
            
            pred_noise = model(x, torch.tensor([t], device=device).repeat(n_samples))
            x = (x - (1 - alpha_t)/torch.sqrt(1 - alpha_bar_t) * pred_noise) / torch.sqrt(alpha_t)
            x += torch.sqrt(beta_t) * z
        
        return x.clamp(-1, 1)

# 4. 训练和保存
def train_and_save():
    device = 'cuda' if torch.cuda.is_available() else 'cpu'
    train_loader = prepare_data()
    model = UNet().to(device)
    diffusion = Diffusion(device=device)
    optimizer = optim.Adam(model.parameters(), lr=2e-4)
    
    # 不再保存训练过程中的图片
    for epoch in range(20):
        total_loss = 0
        for batch, _ in train_loader:
            batch = batch.to(device)
            loss = diffusion.train_step(model, batch, optimizer)
            total_loss += loss
        
        print(f'Epoch {epoch+1}, Loss: {total_loss/len(train_loader):.4f}')
    # 保存模型参数
    torch.save(model.state_dict(), 'mnist_diffusion.pth')

def infer_and_show(n_samples=8, model_path=None, device=None):
    """
    推理生成数字图片并保存（不显示）。
    """
    import matplotlib.pyplot as plt

    if device is None:
        device = 'cuda' if torch.cuda.is_available() else 'cpu'
    model = UNet().to(device)
    diffusion = Diffusion(device=device)

    # 默认加载训练好的模型参数
    if model_path is None:
        model_path = 'mnist_diffusion.pth'
    if os.path.exists(model_path):
        model.load_state_dict(torch.load(model_path, map_location=device))
        print(f"Loaded model weights from {model_path}")

    samples = diffusion.sample(model, n_samples, device)
    samples = (samples + 1) / 2.0  # [-1,1] -> [0,1]

    # 只保存生成的图片，不显示
    os.makedirs('inference_samples', exist_ok=True)
    for i, img in enumerate(samples):
        plt.imsave(f'inference_samples/infer_sample{i}.png', img[0].cpu().numpy(), cmap='gray')

if __name__ == '__main__':
    train_and_save()
    # 推理并显示生成的数字图片
    infer_and_show(n_samples=8)
